'''
土壤水分降尺度算法实现
'''
import os

import numpy as np
from osgeo import gdal, osr
from scipy.optimize import curve_fit


# 读取tif文件
def read_tif(file_path):
    dataset = gdal.Open(file_path)
    data = dataset.ReadAsArray()
    geotransform = dataset.GetGeoTransform()
    return data, geotransform


# 写入tif文件
def write_tif(file_path, data, geotransform, nodata, dataType):
    gdal_type = ''
    if dataType == 'int16':
        gdal_type = gdal.GDT_Int16
    elif dataType == 'int32':
        gdal_type = gdal.GDT_Int32
    elif dataType == 'float32':
        gdal_type = gdal.GDT_Float32
    elif dataType == 'float64':
        gdal_type = gdal.GDT_Float64
    elif dataType == 'byte':
        gdal_type = gdal.GDT_Byte
    elif dataType == 'uint16':
        gdal_type = gdal.GDT_UInt16
    elif dataType == 'uint32':
        gdal_type = gdal.GDT_UInt32
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    dataset.SetGeoTransform(geotransform)
    # 定义投影
    prj = osr.SpatialReference()
    prj.ImportFromEPSG(4326)
    dataset.SetProjection(prj.ExportToWkt())
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    del dataset
    print("写入成功：{}".format(file_path))


# 根据像元行列号计算经纬度
def pixel2coord(geo_transform, x, y):
    x_origin, pixel_width, x_rotation, y_origin, y_rotation, pixel_height = geo_transform
    x_geo = x_origin + pixel_width * x
    y_geo = y_origin + pixel_height * y
    return x_geo, y_geo


# 根据经纬度计算像元行列号
def coord2pixel(geo_transform, x_geo, y_geo):
    x = int((x_geo - geo_transform[0]) / geo_transform[1])
    y = int((y_geo - geo_transform[3]) / geo_transform[5])
    return x, y


# 获取低分辨率像元对应所有高分辨率像元值
def get_high_res_pixels(high_res_geo_transform, high_res_data, low_res_geo_transform, x_low_res_pixel, y_low_res_pixel, low_res_to_high_res_ratio):
    # 计算低分辨率影像像元在高分辨率影像中的范围
    x_low_res_geo, y_low_res_geo = pixel2coord(low_res_geo_transform, x_low_res_pixel, y_low_res_pixel)
    x_low_res_geo = x_low_res_geo + low_res_geo_transform[1] / 2  # 左右偏移
    y_low_res_geo = y_low_res_geo + low_res_geo_transform[5] / 2  # 上下偏移

    # 计算行列号范围
    high_res_x, high_res_y = coord2pixel(high_res_geo_transform, x_low_res_geo, y_low_res_geo)

    high_res_x_min = int(high_res_x - low_res_to_high_res_ratio / 2)
    high_res_x_max = int(high_res_x + low_res_to_high_res_ratio / 2)
    high_res_y_min = int(high_res_y - low_res_to_high_res_ratio / 2)
    high_res_y_max = int(high_res_y + low_res_to_high_res_ratio / 2)

    # 排除边缘像元的影响
    if high_res_x_min < 0 and high_res_x_max > 0:
        high_res_x_min = 0
    if high_res_x_min < high_res_data.shape[1] < high_res_x_max:
        high_res_x_max = high_res_data.shape[1]
    if high_res_y_min < 0 and high_res_y_max > 0:
        high_res_y_min = 0
    if high_res_y_min < high_res_data.shape[0] < high_res_y_max:
        high_res_y_max = high_res_data.shape[0]

    # 获取范围内的像元值
    if 0 <= high_res_x_min <= high_res_data.shape[1] and 0 <= high_res_y_min <= high_res_data.shape[0]:
        # print(high_res_x_min, high_res_y_min, high_res_x_max, high_res_y_max)
        high_res_pixels = high_res_data[high_res_y_min:high_res_y_max, high_res_x_min:high_res_x_max]
        return high_res_pixels.astype(np.float32)
    else:
        return []


# SM和ATI/EVI像元匹配
def pixel_match(ati_path, evi_path, sm_path, snow_path, water_path):
    # 读取高分辨率和低分辨率的影像数据
    ati_data, ati_geotrans = read_tif(ati_path)
    evi_data, evi_geotrans = read_tif(evi_path)
    sm_data, sm_geotrans = read_tif(sm_path)
    snow_dataset = gdal.Open(snow_path)
    water_dataset = gdal.Open(water_path)
    snow_data = snow_dataset.ReadAsArray(xsize=ati_data.shape[0], ysize=ati_data.shape[1])
    water_data = water_dataset.ReadAsArray(xsize=ati_data.shape[0], ysize=ati_data.shape[1])

    # todo 雪和水体掩膜
    mask = (snow_data != 0) & (water_data != 1)
    # ati_data.dtype = np.float
    # evi_data.dtype = np.float
    ati_data = np.where(mask, np.nan, ati_data)
    evi_data = np.where(mask, np.nan, evi_data)

    # print(ati_data.shape)
    # print(evi_data.shape)
    # print(sm_data.shape)
    # 获取像元比例
    low_res_to_high_res_ratio = int(sm_geotrans[1] / ati_geotrans[1])
    # 预设两个列表存每个cci像元对应的ati和evi数据
    ati_in_sm = []
    evi_in_sm = []
    # print(ati_in_sm)
    # 遍历cci每个像元，获取该像元中的所有ati和evi数据（一个二维数组）
    for y in range(0, sm_data.shape[0]):
        for x in range(0, sm_data.shape[1]):
            # print(get_high_res_pixels(ati_geotrans, ati_data, sm_geotrans, x, y, low_res_to_high_res_ratio))
            # print(y * sm_data.shape[0] + x)
            # ati_in_sm[y * sm_data.shape[0] + x] = get_high_res_pixels(ati_geotrans, ati_data, sm_geotrans, x, y, low_res_to_high_res_ratio)
            # evi_in_sm[y * sm_data.shape[0] + x] = get_high_res_pixels(evi_geotrans, evi_data, sm_geotrans, x, y, low_res_to_high_res_ratio)
            ati_in_sm.append(get_high_res_pixels(ati_geotrans, ati_data, sm_geotrans, x, y, low_res_to_high_res_ratio))
            evi_in_sm.append(get_high_res_pixels(evi_geotrans, evi_data, sm_geotrans, x, y, low_res_to_high_res_ratio))
    return np.asarray(ati_in_sm, dtype=np.float32), np.asarray(evi_in_sm, dtype=np.float32), sm_data.flatten()


# 多元线性拟合，输入evi、ati和sm，拟合参数a~f
def line_fit(ati_path, evi_path, sm_path, snow_path, water_path, fit_model):
    # 获取ati、evi和sm的一维数据格式，注意的是ati和evi每个元素是一个数组（这个数组对应sm一个像元的所有数据）
    ati, evi, sm = pixel_match(ati_path, evi_path, sm_path, snow_path, water_path)
    # print(ati.shape)
    # print(evi.shape)
    # print(sm.shape)
    # print(ati)
    # todo 剔除异常值，cci数据-9999不要
    # print(sm != -9999)
    ati_filter = ati[sm != -9999]
    evi_filter = evi[sm != -9999]
    sm_filter = sm[sm != -9999]
    # print(ati_filter.shape)
    # print(evi_filter.shape)
    # print(sm_filter.shape)

    # print(ati_filter[0])
    # print(evi_filter[0])
    # todo 剔除ati和evi中的异常值
    # 方法一：替换成nan
    ati_filter_new = [np.where(arr == -9999., np.nan, arr) for arr in ati_filter]
    ati_filter_new = [np.where(arr == np.inf, np.nan, arr) for arr in ati_filter_new]
    evi_filter_new = [np.where(arr == -3000., np.nan, arr) for arr in evi_filter]
    evi_filter_new = [np.where(arr == np.inf, np.nan, arr) for arr in evi_filter_new]

    # 将每个元素展平并合并为二维数组（为了用于curve_fit拟合用，curve_fit只能输入一个自变量，所以需要将两个变量组合成一个来传输）
    X = np.vstack((np.concatenate([arr.flatten() for arr in ati_filter_new]), np.concatenate([arr.flatten() for arr in evi_filter_new]))).T
    popt, pcov = curve_fit(fit_model, X, sm_filter, maxfev=1000000, check_finite=False)
    return popt, pcov


# 拟合公式
def fit_formula(X, a, b, c, d, e, f):
    ati, evi = X.T  # 转置以匹配合并的数组
    # 经验公式
    sm = a * np.nanmean(np.power(d, (evi + e))) + b * np.nanmean(np.log(ati + c) * np.power(d, (evi + e))) + f
    # sm = a * np.nanmean(np.power(d, (evi + e))) + b * np.nanmean(np.log(ati + c + 1e-5) * np.power(d, (evi + e))) + f  # 1e-5是为了排除当数值很小时，取对数后结果趋于负无穷大
    return sm


# 降尺度公式
def down_scaling(X, a, b, c, d, e, f):
    ati, evi = X.T  # 转置以匹配合并的数组
    sm = (a + b * np.log(ati + c)) * np.power(d, (evi + e)) + f
    return sm


def main(ati_path, evi_path, sm_path, snow_path, water_path, output_path):
    # 运行多元拟合函数
    popt, pcov = line_fit(ati_path, evi_path, sm_path, snow_path, water_path, fit_formula)
    a, b, c, d, e, f = popt[0], popt[1], popt[2], popt[3], popt[4], popt[5]
    print("拟合参数:")
    print(f"a: {a}")
    print(f"b: {b}")
    print(f"c: {c}")
    print(f"d: {d}")
    print(f"e: {e}")
    print(f"f: {f}")

    ati_data, ati_geotrans = read_tif(ati_path)
    evi_data, evi_geotrans = read_tif(evi_path)
    snow_dataset = gdal.Open(snow_path)
    water_dataset = gdal.Open(water_path)
    snow_data = snow_dataset.ReadAsArray(xsize=ati_data.shape[0], ysize=ati_data.shape[1])
    water_data = water_dataset.ReadAsArray(xsize=ati_data.shape[0], ysize=ati_data.shape[1])

    # todo 雪和水体掩膜
    mask = (snow_data != 0) & (water_data != 1)
    ati_data = np.where(mask, np.nan, ati_data)
    evi_data = np.where(mask, np.nan, evi_data)

    X = np.vstack((ati_data.flatten(), evi_data.flatten())).T
    # 调用降尺度公式
    sm_high_res = down_scaling(X, a, b, c, d, e, f)
    sm_high_res = sm_high_res.reshape(ati_data.shape)
    write_tif(output_path, sm_high_res, ati_geotrans, np.nan, 'float32')


if __name__ == '__main__':
    # 指定输入文件路径，假设都转为tif格式了
    # ATI_PATH = r"G:\遥感图像处理\土壤水分降尺度\207-ati.tif"
    # EVI_PATH = r"G:\遥感图像处理\土壤水分降尺度\207-evi.tif"
    # SM_PATH = r"G:\遥感图像处理\土壤水分降尺度\207-cci.tif"
    # OUTPUT_PATH = r"G:\遥感图像处理\土壤水分降尺度\cci_high_207.tif"
    # main(ATI_PATH, EVI_PATH, SM_PATH, OUTPUT_PATH)

    # 批量
    ati_dir = r"G:\test\process_result\ATI_tif"
    evi_dir = r"G:\test\process_result\EVI_tif\EVI_quadratic_2013"
    cci_dir = r"G:\test\process_result\ESACCI_tif\clip"
    snow_dir = r"G:\test\process_result\MCD_tif\snow"
    water_dir = r"G:\test\process_result\MCD_tif\water"
    output_dir = r"G:\test\process_result\downscale_result"
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    ati_file_list = os.listdir(ati_dir)
    cci_path_list = [os.path.join(cci_dir, cci) for cci in os.listdir(cci_dir)]
    for i, ati_file in enumerate(ati_file_list):
        print(ati_file)
        ati_path = os.path.join(ati_dir, ati_file)
        evi_path = os.path.join(evi_dir, ati_file)
        cci_path = cci_path_list[i]
        snow_path = os.path.join(snow_dir, ati_file)
        water_path = os.path.join(water_dir, ati_file)
        output_path = os.path.join(output_dir, ati_file)
        main(ati_path, evi_path, cci_path, snow_path, water_path, output_path)
        print("=" * 30)
